#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
@author coldwind
"""

import tensorflow  as tf


class TextBIRNN:
    def __init__(self,
                 model_type,
                 sequence_length,
                 number_class,
                 word_embedding,
                 embedding_size,
                 vocab_size,
                 number_layers,
                 number_units,
                 model='lstm',
                 l2_reg_lambda=0.0):
        self.input_x = tf.placeholder(shape=[None, sequence_length], dtype=tf.int32, name='input_x')
        self.input_y = tf.placeholder(shape=[None, number_class], dtype=tf.float32, name='input_y')
        self.sequence_real_length = tf.placeholder(shape=[None], dtype=tf.int32, name='sequence_real_length')
        self.dropout_keep_prob = tf.placeholder(dtype=tf.float32, name='dropout_keep_prob')
        self.learning_rate = tf.placeholder(dtype=tf.float32, name='learning_rate')

        self.l2_loss = tf.constant(0.0, dtype=tf.float32)

        # embedding
        with tf.device('/cpu:0'), tf.name_scope('embedding'):
            if word_embedding:
                self.Embedding = tf.Variable(initial_value=word_embedding, dtype=tf.float32, name='word_embedding')
            else:
                self.Embedding = tf.get_variable(
                    initializer=tf.random_uniform(shape=[vocab_size, embedding_size], minval=-1, maxval=1),
                    dtype=tf.float32, name='word_embedding')

            self.embedded_words = tf.nn.embedding_lookup(self.Embedding, self.input_x, name='embedded_words')

        # 构建rnn
        with tf.name_scope(model):
            if 'lstm' == model:
                cell_func = tf.nn.rnn_cell.LSTMCell
            elif 'gru' == model:
                cell_func = tf.nn.rnn_cell.GRUCell
            elif 'rnn' == model:
                cell_func = tf.nn.rnn_cell.BasicRNNCell

            def get_cell(num_units):
                cell = cell_func(num_units=num_units, state_is_tuple=True)
                output = tf.nn.rnn_cell.DropoutWrapper(cell=cell, output_keep_prob=self.dropout_keep_prob)

                return output

            def get_bi_rnn(inputs, scope, num_units):
                # https://www.tensorflow.org/versions/r1.12/api_docs/python/tf/nn/bidirectional_dynamic_rnn
                # bidirectional_dynamic_rnn(
                #     cell_fw,  # 前向RNN
                #     cell_bw,  # 后向RNN
                #     inputs,  # 输入
                #     sequence_length=None,  # 输入序列的实际长度（可选，默认为输入序列的最大长度）
                #     initial_state_fw=None,  # 前向的初始化状态（可选）
                #     initial_state_bw=None,  # 后向的初始化状态（可选）
                #     dtype=None,  # 初始化和输出的数据类型（可选）
                #     parallel_iterations=None,
                #     swap_memory=False,
                #     time_major=False,
                #     # 决定了输入输出tensor的格式：如果为true, 向量的形状必须为 `[max_time, batch_size, depth]`.
                #     # 如果为false, tensor的形状必须为`[batch_size, max_time, depth]`.
                #     scope=None
                # )
                cell_fw = get_cell(num_units)
                cell_bw = get_cell(num_units)
                # outputs:
                #     (
                #       [batch_size, max_seq_length, number_units],
                #       [batch_size, max_seq_length, number_units]
                #     )
                outputs, _ = tf.nn.bidirectional_dynamic_rnn(cell_fw=cell_fw, cell_bw=cell_bw,
                                                             inputs=inputs,
                                                             sequence_length=self.sequence_real_length,
                                                             dtype=tf.float32, scope=scope)

                # [batch_size, max_seq_length, number_units * 2],
                outputs = tf.concat(outputs, axis=2)

                return outputs

            output = self.embedded_words
            for i in range(number_layers):
                output = get_bi_rnn(inputs=output, scope='layer-%d' % i, num_units=number_units)

            self.output = tf.reduce_mean(output, axis=1)

        # dropout
        with tf.name_scope('dropout'):
            self.h_dropout = tf.nn.dropout(self.output, keep_prob=self.dropout_keep_prob, name='dropout')

        # full-connection
        with tf.name_scope('output'):
            in_dim = self.h_dropout.shape[-1].value
            w = tf.get_variable(shape=[in_dim, number_class], dtype=tf.float32,
                                initializer=tf.contrib.layers.xavier_initializer(), name='W')
            b = tf.Variable(tf.constant(value=0.1, shape=[number_class]), dtype=tf.float32, name='b')

            self.l2_loss += tf.nn.l2_loss(w)
            self.l2_loss += tf.nn.l2_loss(b)

            self.scores = tf.nn.xw_plus_b(x=self.h_dropout, weights=w, biases=b, name='scores')
            if 'cls' == model_type:
                self.predictions = tf.argmax(self.scores, axis=1, name='predictions')
            elif 'reg' == model_type:
                self.predictions = tf.reduce_max(self.scores, axis=1, name='predictions')
                self.predictions = tf.expand_dims(self.predictions, -1)

        with tf.name_scope('loss'):
            if 'cls' == model_type:
                losses = tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.input_y, logits=self.scores)
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda
            elif 'reg' == model_type:
                losses = tf.sqrt(tf.losses.mean_squared_error(labels=self.input_y, predictions=self.predictions))
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda

        with tf.name_scope('accuracy'):
            if 'cls' == model_type:
                correct_predictions = tf.equal(tf.argmax(self.input_y, axis=1), self.predictions)
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, dtype=tf.float32), name='accuracy')
            elif 'reg' == model_type:
                self.accuracy = tf.constant(0.0, name='accuracy')


if __name__ == '__main__':
    text_birnn = TextBIRNN(model_type='cls', sequence_length=10, number_class=2, word_embedding=None,
                           embedding_size=128, vocab_size=100, number_layers=2, number_units=64, model='lstm',
                           l2_reg_lambda=0.0001)
